#include <opencv2/opencv.hpp>
#include <iostream>
#include <string>
#include <vector>

class PolygonDetector {
private:
    cv::Mat srcImage;
    cv::Mat resultImage;
    std::vector<cv::Point2f> imagePoints; // 图像点（4个，匹配4边形灯条）
    std::vector<cv::Point3f> objectPoints; // 世界点（4个，按灯条实际尺寸初始化）

public:
    PolygonDetector(const std::string& imagePath) {
        srcImage = cv::imread(imagePath);
        if (srcImage.empty()) {
            std::cerr << "Failed to load image: " << imagePath << std::endl;
        }
        // 关键修改1：世界点改为4个（示例为矩形灯条，需按实际4边形尺寸调整！单位：mm）
        objectPoints = {
            cv::Point3f(0, 0, 0),         // 4边形顶点1（世界坐标系原点）
            cv::Point3f(200, 0, 0),       // 顶点2（假设底边水平长度200mm）
            cv::Point3f(200, 150, 0),     // 顶点3（假设垂直高度150mm）
            cv::Point3f(0, 150, 0)        // 顶点4（闭合4边形）
        };
    }

    bool detectPolygon() {
        if (srcImage.empty()) {
            return false;
        }
        resultImage = srcImage.clone();

        // 图像预处理（保留调试用二值图）
        cv::Mat grayImage, blurImage, binaryImage;
        cv::cvtColor(srcImage, grayImage, cv::COLOR_BGR2GRAY);
        cv::GaussianBlur(grayImage, blurImage, cv::Size(5, 5), 0);
        cv::threshold(blurImage, binaryImage, 0, 255, cv::THRESH_BINARY_INV | cv::THRESH_OTSU);
        cv::imwrite("../../images/binary_debug.jpg", binaryImage);

        // 查找轮廓
        std::vector<std::vector<cv::Point>> contours;
        std::vector<cv::Vec4i> hierarchy;
        cv::findContours(binaryImage, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        std::vector<cv::Point2f> validImagePoints;
        std::cout << "\n=== [调试信息] 轮廓识别结果 ===" << std::endl;
        std::cout << "共检测到轮廓数量：" << contours.size() << std::endl;

        for (size_t i = 0; i < contours.size(); i++) {
            // 面积过滤阈值（保持30，适配当前轮廓面积88995）
            double area = cv::contourArea(contours[i]);
            if (area < 30) {
                std::cout << "轮廓" << i+1 << "：面积=" << area << "（小于30，过滤）" << std::endl;
                continue;
            }

            // 多边形逼近精度（保持0.025，当前已稳定识别4边形）
            std::vector<cv::Point> approx;
            double epsilon = cv::arcLength(contours[i], true) * 0.025;
            cv::approxPolyDP(contours[i], approx, epsilon, true);

            std::cout << "轮廓" << i+1 << "：面积=" << area << "，逼近后边数=" << approx.size() << std::endl;
            // 关键修改2：仅保留4边形轮廓（匹配世界点数量）
            if (approx.size() == 4) {
                // 颜色判断与增强标注
                cv::Rect rect = cv::boundingRect(contours[i]);
                cv::Mat roi = srcImage(rect);
                cv::Scalar meanColor = cv::mean(roi);
                std::string colorName = getColorName(meanColor);

                cv::drawContours(resultImage, contours, i, cv::Scalar(0, 0, 255), 8);
                std::string info = colorName + ", Quadrilateral (4 sides, Area:" + std::to_string((int)area) + ")";
                cv::putText(resultImage, info, cv::Point(rect.x, rect.y - 10),
                            cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 0, 255), 3);

                // 保存4个顶点作为有效图像点（仅保留第一个符合条件的4边形）
                if (validImagePoints.empty()) {
                    for (const auto& point : approx) {
                        validImagePoints.push_back(cv::Point2f(point.x, point.y));
                    }
                    std::cout << "轮廓" << i+1 << "：已选为有效多边形（保存4个顶点）" << std::endl;
                }
            } else {
                std::cout << "轮廓" << i+1 << "：边数≠4（过滤）" << std::endl;
            }
        }

        imagePoints = validImagePoints;
        std::cout << "=== [调试信息] 有效图像点统计 ===" << std::endl;
        std::cout << "最终保存的图像点数量：" << imagePoints.size() << std::endl;
        return true;
    }

    // 颜色判断辅助函数（不变）
    std::string getColorName(const cv::Scalar& color) {
        double b = color[0], g = color[1], r = color[2];
        if (r > g && r > b) return "Red";
        else if (g > r && g > b) return "Green";
        else if (b > r && b > g) return "Blue";
        else return "Other";
    }

    void solvePose() {
        std::cout << "\n=== [调试信息] 位姿解算准备 ===" << std::endl;
        // 关键修改3：点数判断改为4个（匹配4边形）
        if (imagePoints.size() != 4 || objectPoints.size() != 4) {
            std::cerr << "[Error] 位姿解算失败：图像点/世界点数量不匹配！" << std::endl;
            std::cerr << "- 有效图像点数量：" << imagePoints.size() << "（需4个）" << std::endl;
            std::cerr << "- 预设世界点数量：" << objectPoints.size() << "（已满足4个）" << std::endl;
            return;
        }

        // 相机内参（需替换为实际标定值，示例适配640x480分辨率）
        cv::Mat cameraMatrix = (cv::Mat_<double>(3, 3) <<
            800, 0, 320,   // fx=800, cx=320
            0, 800, 240,   // fy=800, cy=240
            0, 0, 1);
        cv::Mat distCoeffs = cv::Mat::zeros(5, 1, CV_64F); // 畸变系数（未标定则设为0）

        cv::Mat rvec, tvec;
        // 关键修改4：4点解算用默认方法（无需指定flags，稳定性更高）
        bool solveSuccess = cv::solvePnP(
            objectPoints, imagePoints, cameraMatrix, distCoeffs,
            rvec, tvec
        );

        if (solveSuccess) {
            cv::Mat rotationMatrix;
            cv::Rodrigues(rvec, rotationMatrix); // 旋转向量转3x3旋转矩阵
            std::cout << "\n=== 灯条位姿解算成功 ===" << std::endl;
            std::cout << "1. 平移向量 (Translation Vector) [单位：mm]：" << std::endl << tvec << std::endl;
            std::cout << "2. 旋转矩阵 (Rotation Matrix)：" << std::endl << rotationMatrix << std::endl;
        } else {
            std::cerr << "[Error] solvePnP解算失败！可能原因：图像点与世界点顺序不匹配、相机内参错误" << std::endl;
            std::cerr << "[提示] 可检查objectPoints的4个顶点顺序，需与imagePoints的顶点顺序一致（如均为顺时针/逆时针）" << std::endl;
        }
    }

    cv::Mat getResultImage() const {
        return resultImage;
    }
};

int main() {
    // 图片路径（适配可执行文件在build/bin目录，指向项目根目录images文件夹）
    std::string imagePath = "../../images/robomaster_sample.jpg";

    PolygonDetector detector(imagePath);
    if (!detector.detectPolygon()) {
        std::cerr << "图像识别初始化失败！" << std::endl;
        return -1;
    }

    detector.solvePose();

    // 显示并保存调试版结果图
    cv::imshow("Detection & Pose Result (4-Point Version)", detector.getResultImage());
    cv::imwrite("../../images/result_4point.jpg", detector.getResultImage());

    cv::waitKey(0);
    cv::destroyAllWindows();
    return 0;
}